# app_fixed.py
from fastapi import FastAPI, WebSocket, WebSocketDisconnect, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from typing import Dict, List, Optional
import json
import asyncio
from datetime import datetime, timedelta
import uuid
import logging
from enum import Enum
from contextlib import asynccontextmanager

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)


class DeviceStatus(BaseModel):
    """设备状态模型"""
    deviceId: str
    versionName: str
    url: str
    code: str
    connected_at: Optional[datetime] = None
    last_heartbeat: Optional[datetime] = None
    is_online: bool = False




class MessageType(str, Enum):
    ERROR = "error"
    DEVICE_REGISTER = "device_register"
    HEARTBEAT = "heartbeat"
    COMMAND = "command"
    STATUS_UPDATE = "status_update"
    BROADCAST = "broadcast"
    DATA_REPORT = "data_report"


class WebSocketMessage(BaseModel):
    type: MessageType
    data: Dict
    timestamp: Optional[datetime] = None
    target_device: Optional[str] = None


class ConnectionManager:
    def __init__(self):
        self.active_connections: Dict[str, WebSocket] = {}
        self.device_statuses: Dict[str, DeviceStatus] = {}
        self.scheduled_tasks: Dict[str, asyncio.Task] = {}
        self.health_check_task: Optional[asyncio.Task] = None

    async def connect(self, websocket: WebSocket, device_id: str):
        """连接设备 - 修复版本"""
        try:
            await websocket.accept()
            self.active_connections[device_id] = websocket
            logger.info(f"设备 {device_id} WebSocket 连接已接受")
            return True
        except Exception as e:
            logger.error(f"接受设备 {device_id} 连接失败: {e}")
            return False

    def disconnect(self, device_id: str):
        """断开设备连接"""
        if device_id in self.active_connections:
            try:
                # 注意：这里不调用 websocket.close()，因为连接可能已经关闭
                del self.active_connections[device_id]
                logger.info(f"设备 {device_id} 从活跃连接中移除")
            except Exception as e:
                logger.error(f"移除设备{device_id}连接时出错: {e}")

                # 更新设备状态
                if device_id in self.device_statuses:
                    self.device_statuses[device_id].is_online = False
                self.device_statuses[device_id].last_heartbeat = datetime.now()

                # 取消定时任务
                if device_id in self.scheduled_tasks:
                    try:
                        self.scheduled_tasks[device_id].cancel()
                        del self.scheduled_tasks[device_id]
                    except Exception as e:
                        logger.error(f"取消设备 {device_id} 定时任务时出错: {e}")

    async def register_device(self, device_status: DeviceStatus, websocket: WebSocket):
        """注册设备"""
        # try:
        # device_status.connected_at = datetime.now()
        # device_status.last_heartbeat = datetime.now()
        device_status.is_online = True

        self.device_statuses[device_status.deviceId] = device_status
        return True
        # 先连接再注册
        # if await self.connect(websocket, device_status.deviceId):
        #     logger.info(f"设备注册成功: {device_status.deviceId}")
        #     self._start_scheduled_messages(device_status.deviceId)
        #     return True
        # else:
        #     logger.error(f"设备注册失败: {device_status.deviceId} - 连接失败")
        #     return False

        # except Exception as e:
        #     logger.error(f"设备注册异常 {device_status.deviceId}: {e}")
        #     return False

    def _start_scheduled_messages(self, device_id: str):
        """启动定时消息"""
        if device_id in self.scheduled_tasks:
            try:
                self.scheduled_tasks[device_id].cancel()
            except:
                pass

        task = asyncio.create_task(self._send_scheduled_messages(device_id))
        self.scheduled_tasks[device_id] = task

    async def _send_scheduled_messages(self, device_id: str):
        """发送定时消息"""
        try:
            while True:
                await asyncio.sleep(30)

                if device_id in self.active_connections and device_id in self.device_statuses:
                    if self.device_statuses[device_id].is_online:
                        message = WebSocketMessage(
                            type=MessageType.COMMAND,
                            data={
                                "command": "status_query",
                                "message": "定时状态查询",
                                "scheduled": True
                            },
                            timestamp=datetime.now()
                        )
                        await self.send_personal_message(json.dumps(message.dict()), device_id)
                        logger.debug(f"向设备 {device_id} 发送定时状态查询")

        except asyncio.CancelledError:
            logger.info(f"设备 {device_id} 定时任务已取消")
        except Exception as e:
            logger.error(f"设备 {device_id} 定时消息错误: {e}")

    async def send_personal_message(self, message: str, device_id: str):
        """发送个人消息"""
        if device_id in self.active_connections:
            try:
                await self.active_connections[device_id].send_text(message)
                return True
            except Exception as e:
                logger.error(f"向设备 {device_id} 发送消息失败: {e}")
                self.disconnect(device_id)
                return False
        return False

    async def broadcast(self, message: str):
        """广播消息"""
        success_count = 0
        failed_devices = []

        for device_id, connection in list(self.active_connections.items()):
            try:
                await connection.send_text(message)
                success_count += 1
            except Exception as e:
                logger.error(f"广播到设备 {device_id} 失败: {e}")
                failed_devices.append(device_id)

        # 清理失败的连接
        for device_id in failed_devices:
            self.disconnect(device_id)

        return success_count

    async def update_heartbeat(self, device_id: str):
        """更新心跳"""
        if device_id in self.device_statuses:
            self.device_statuses[device_id].last_heartbeat = datetime.now()
            self.device_statuses[device_id].is_online = True

    def get_online_devices(self) -> List[DeviceStatus]:
        """获取在线设备"""
        return [status for status in self.device_statuses.values() if status.is_online]

    def get_device_status(self, device_id: str) -> Optional[DeviceStatus]:
        """获取设备状态"""
        return self.device_statuses.get(device_id)

    async def send_command_to_device(self, device_id: str, command: Dict):
        """发送命令到设备"""
        message = WebSocketMessage(
            type=MessageType.COMMAND,
            data=command,
            timestamp=datetime.now()
        )
        return await self.send_personal_message(json.dumps(message.dict()), device_id)

    async def start_health_check(self):
        """启动健康检查"""
        self.health_check_task = asyncio.create_task(self._health_check())

    async def stop_health_check(self):
        """停止健康检查"""
        if self.health_check_task:
            self.health_check_task.cancel()
            try:
                await self.health_check_task
            except asyncio.CancelledError:
                pass

    async def _health_check(self):
        """健康检查"""
        try:
            while True:
                await asyncio.sleep(60)
                current_time = datetime.now()
                offline_count = 0

                for device_id, status in self.device_statuses.items():
                    if status.last_heartbeat and status.is_online:
                        if current_time - status.last_heartbeat > timedelta(minutes=2):
                            status.is_online = False
                            offline_count += 1
                            logger.warning(f"设备 {device_id} 心跳超时，标记为离线")

                online_count = len(self.get_online_devices())
                logger.info(f"健康检查: {online_count} 在线, {offline_count} 离线")

        except asyncio.CancelledError:
            logger.info("健康检查任务已停止")
        except Exception as e:
            logger.error(f"健康检查错误: {e}")


# 创建管理器
manager = ConnectionManager()


@asynccontextmanager
async def lifespan(app: FastAPI):
    # 启动
    await manager.start_health_check()
    logger.info("🚀 设备管理服务器启动完成")
    yield
    # 关闭
    await manager.stop_health_check()
    logger.info("🛑 设备管理服务器已关闭")


app = FastAPI(title="设备管理 WebSocket 服务器", lifespan=lifespan)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


@app.websocket("/ws/device/{device_id}")
async def websocket_endpoint(websocket: WebSocket, device_id: str):
    """WebSocket 端点 - 修复版本"""
    logger.info(f"接收到设备连接请求: {device_id}")

    try:
        # 先建立 WebSocket 连接
        await websocket.accept()
        logger.info(f"设备 {device_id} WebSocket 连接已建立")

        # 等待设备发送注册消息
        while True:
            data = await websocket.receive_text()
            print(data)
            message_dict = json.loads(data)

            try:
                message = WebSocketMessage(**message_dict)
            except Exception as e:
                logger.error(f"设备 {device_id} 消息解析失败: {e}")
                error_msg = WebSocketMessage(
                    type=MessageType.ERROR,
                    data={"error": "INVALID_MESSAGE", "description": str(e)},
                    timestamp=datetime.now()
                )
                await websocket.send_text(json.dumps(error_msg.dict()))
                continue

            # 处理注册消息
            if message.type == MessageType.DEVICE_REGISTER:
                print(message.type)
                device_status = DeviceStatus(
                    deviceId=device_id,
                    versionName=message.data.get("versionName"),
                    url=message.data.get("url"),
                    code=message.data.get("code")
                )
                print(str(device_status))
                if await manager.register_device(device_status, websocket):
                    response = WebSocketMessage(
                        type=MessageType.STATUS_UPDATE,
                        data={"status": "registered", "message": "设备注册成功"},
                        timestamp=datetime.now()
                    )

                    message_dict = response.dict()
                    if message_dict.get('timestamp'):
                        message_dict['timestamp'] = message_dict['timestamp'].isoformat()

                    await websocket.send_text(json.dumps(message_dict))
                else:
                    error_msg = WebSocketMessage(
                        type=MessageType.ERROR,
                        data={"error": "REGISTRATION_FAILED", "description": "设备注册失败"},
                        timestamp=datetime.now()
                    )
                    await websocket.send_text(json.dumps(error_msg.dict()))
                    break

            elif message.type == MessageType.HEARTBEAT:
                await manager.update_heartbeat(device_id)
                response = WebSocketMessage(
                    type=MessageType.HEARTBEAT,
                    data={"status": "ok", "message": "心跳接收成功"},
                    timestamp=datetime.now()
                )
                await websocket.send_text(json.dumps(response.dict()))

            elif message.type == MessageType.STATUS_UPDATE:
                logger.info(f"设备 {device_id} 状态更新: {message.data}")

            elif message.type == MessageType.DATA_REPORT:

                logger.info(f"设备 {device_id} 数据：{message_dict['data']['content']}")
                logger.info(f"设备 {device_id} 数据报告: {len(message.data.get('data_points', []))} 个数据点")

            else:
                logger.info(f"设备 {device_id} 发送消息: {message.type}")

    except WebSocketDisconnect:
        logger.info(f"设备 {device_id} 断开连接")
        manager.disconnect(device_id)
    # except Exception as e:
    #     logger.error(f"设备 {device_id} 处理异常: {e}")
    #     manager.disconnect(device_id)


# REST API 保持不变...
@app.get("/")
async def root():
    return {"message": "设备管理 WebSocket 服务器", "status": "running"}


@app.get("/devices/online")
async def get_online_devices():
    return manager.get_online_devices()


@app.get("/devices/{device_id}")
async def get_device_status(device_id: str):
    status = manager.get_device_status(device_id)
    if not status:
        raise HTTPException(status_code=404, detail="设备未找到")
    return status


@app.post("/devices/{device_id}/command")
async def send_command_to_device(device_id: str, command: Dict):
    if device_id not in manager.active_connections:
        raise HTTPException(status_code=404, detail="设备不在线")

    success = await manager.send_command_to_device(device_id, command)
    if success:
        return {"message": "命令已发送", "device_id": device_id}
    else:
        raise HTTPException(status_code=500, detail="命令发送失败")


@app.post("/broadcast")
async def broadcast_message(message: Dict):
    ws_message = WebSocketMessage(
        type=MessageType.BROADCAST,
        data=message,
        timestamp=datetime.now()
    )
    success_count = await manager.broadcast(json.dumps(ws_message.dict()))
    return {
        "message": "广播消息已发送",
        "success_count": success_count,
        "total_devices": len(manager.active_connections)
    }


@app.delete("/devices/{device_id}")
async def disconnect_device(device_id: str):
    manager.disconnect(device_id)
    return {"message": f"设备 {device_id} 已断开"}


if __name__ == "__main__":
    import uvicorn

    uvicorn.run("app:app", host="0.0.0.0", port=8000, reload=True)